#!/usr/bin/env python3
# ! <<< Add self
import datetime
import pathlib
# ! >>>
# ! <<<
# from keras.callbacks import *
# from keras.optimizers import *
from tensorflow.keras.callbacks import *
from tensorflow.keras.optimizers import *
# ! >>>

import sys,os
# ! <<< Outdated
import tensorflow as tf
# ! >>>

sys.path.append('..')
sys.path.append('../..')

"""The image format of all models is "channels_last"""
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

from paraClass import *
import Unet as Unet
from diceLoss import *
import dataGenerator as DGen

def ontrain_2D(organ,inputshape,imgpath,labelpath,traintxtpath,testtxtpath,batchsize,modelprefix,sampleparas,normparas,savepath):
    model = Unet.get_2D_Unet(inputshape,sampleparas.nclass,activation='sigmoid',axis=-1)
    # ! <<< Outdated
    # model.compile(optimizer = Adam(lr = 1e-3,clipnorm=1.), loss = dice_coef_loss, metrics = [dice_coef])
    model.compile(optimizer = tf.keras.optimizers.Adam(lr = 1e-3,clipnorm=1.), loss = dice_coef_loss, metrics = [dice_coef])
    # model.compile(optimizer = 'adam', loss = dice_coef_loss, metrics = [dice_coef])
    # ! >>>
    model_savepath = savepath + "/%s-%s-{epoch:02d}-{dice_coef:.4f}-{val_dice_coef:.4f}.hdf5"%(organ,modelprefix)
    model_checkpoint = ModelCheckpoint(model_savepath, monitor='val_loss',verbose=1, save_best_only=True)
    model_reducelr = ReduceLROnPlateau(monitor='val_loss', factor=0.1, patience=5, verbose=0, mode='auto', epsilon=0.000001, cooldown=0, min_lr=0)
    print('Fitting model...')

    # ! <<< Outdated
    # model.fit_generator(DGen.online_batch_generator_2D(imgpath,labelpath,batchsize,sampleparas,normparas,txtpath=traintxtpath), \
    model.fit(DGen.online_batch_generator_2D(imgpath,labelpath,batchsize,sampleparas,normparas,txtpath=traintxtpath), \
    validation_data=DGen.online_batch_generator_2D(imgpath,labelpath,batchsize,sampleparas,normparas,txtpath=testtxtpath), \
    validation_steps = 250, steps_per_epoch = 250, epochs=300, verbose=1, shuffle=True, \
    callbacks=[model_checkpoint,model_reducelr])
    # validation_steps = 100, steps_per_epoch = 1000, epochs=100, verbose=1, shuffle=True, \

if __name__ == '__main__':
    organname = 'GTV-T'

    # imgpath = '/media/shuaiw/DataDisk/S_Wang_Research/P_Liming_RatBrain_Extraction/TrainData/Total_Mouth_Rat_Resampled/Raw'
    # labelpath = '/media/shuaiw/DataDisk/S_Wang_Research/P_Liming_RatBrain_Extraction/TrainData/Total_Mouth_Rat_Resampled/Mask'
    # savepath = '/media/shuaiw/DataDisk/S_Wang_Research/P_Liming_RatBrain_Extraction/ModelZoo'

    # imgpath = '/home/yusongli/_dataset/shidaoai/img/_out/wangqifeng_spacial_cropped_96_224_224_fix_affine_dilated_img_mask/img'
    # labelpath = '/home/yusongli/_dataset/shidaoai/img/_out/wangqifeng_spacial_cropped_96_224_224_fix_affine_dilated_img_mask/mask'
    imgpath = None
    labelpath = None

    # traintxtpath = '/media/shuaiw/DataDisk/S_Wang_Research/P_Liming_RatBrain_Extraction/TrainData/Total_Mouth_Rat_Resampled/RatBrain_train.txt'
    # testtxtpath = '/media/shuaiw/DataDisk/S_Wang_Research/P_Liming_RatBrain_Extraction/TrainData/Total_Mouth_Rat_Resampled/RatBrain_val.txt'

    savepath = f"/home/yusongli/_dataset/shidaoai/img/_out/2dunet/{datetime.datetime.now().strftime('%Y%m%d')}"
    traintxtpath = '/home/yusongli/Documents/shidaoai_new_project/networks/2D-Unet/_conf/train.txt'
    testtxtpath = '/home/yusongli/Documents/shidaoai_new_project/networks/2D-Unet/_conf/val.txt'

    pathlib.Path(savepath).mkdir(exist_ok=True, parents=True)

    modelprefix = '2D_Unet'
    # ? <<< Self parameters
    # inputshape = [128,128,1]
    # inputshape = [26,26,1]
    inputshape = [16, 16, 1]
    # ? >>>
    # ! <<< Self parameters
    # batchsize = 2
    batchsize = 6
    # batchsize = 1
    # ! >>>

    normparas = NormParas()
    normparas.method = "minmax"
    normparas.lmin = 0
    # ! <<< Self parameters
    # normparas.rmax = 2000
    normparas.rmax = 1500
    # ! >>>

    sampleparas = SampleParas()
    # ? <<< Self parameters
    # sampleparas.patchdims = [1, 128, 128]
    # sampleparas.patchlabeldims = [1, 128, 128]
    sampleparas.patchdims = [1, 16, 16]
    sampleparas.patchlabeldims = [1, 16, 16]
    # ? >>>
    sampleparas.imgdim = 2
    sampleparas.nclass = 2
    sampleparas.samplenum = 8

    ontrain_2D(organname,inputshape,imgpath,labelpath,traintxtpath,testtxtpath,batchsize,modelprefix,sampleparas,normparas,savepath)

